import random
import time

import numpy as np
import vrplib as vrp
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties

# 设置中文字体，这里使用SimHei字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False  # 解决负号“-”显示为方块的问题

POPULATION_SIZE = 1000
N_GENERATIONS = 100
MAX_CAPACITY = 100
MUTATION_RATE = 0.5

# vrp.download_instance('A-n33-k5', 'data')
instance = vrp.read_instance(r'data\A-n32-k5.vrp')
demand = instance['demand']
weights = instance['edge_weight']
custom_list = np.arange(1, len(demand))
coord = instance['node_coord'].T
print("demand", demand)
print('weights', weights)
print('custom_list', custom_list)
print('coord', coord)


def solution_routes(x):
    routes, route, route_sum = [], [], 0
    for i in x:
        if route_sum + demand[i] >= MAX_CAPACITY:
            routes.append(route)
            route, route_sum = [], 0
        route.append(i)
        route_sum += demand[i]
    if (len(route) > 0):
        routes.append(route)

    return routes


def routes_cost(routes):
    cost = 0
    for r in routes:
        rlist = [0] + r + [0]
        for i in range(len(rlist) - 1):
            cost += weights[rlist[i]][rlist[i + 1]]
    print(weights)

    return cost


# 适应度函数
def fitness_function(x):
    routes = solution_routes(x)
    cost = routes_cost(routes)
    return 1 / cost


# 选择函数
def selection(population, fitness_values):
    next_population = []
    for _ in range(POPULATION_SIZE):
        parent = random.choices(population, weights=fitness_values)[0]
        next_population.append(parent)

    return next_population


# 变异函数
def mutation(population):
    mut_population = [x.copy() for x in population]
    for i in range(len(mut_population)):
        k = \
            random.choices(np.arange(2, len(demand) - 1),
                           weights=(1 / (np.arange(2, len(demand) - 1))) ** MUTATION_RATE)[0]
        # k代表变异长度
        # print(k)
        idxs = list(set([random.randint(0, len(demand) - 2) for _ in range(k)]))
        mut_population[i][idxs] = np.random.permutation(mut_population[i][idxs])
    all_population = population + mut_population
    fitness_values = [fitness_function(x) for x in all_population]
    p, topk = fitness_values, POPULATION_SIZE
    index_list = sorted(range(len(p)), key=lambda i: p[i], reverse=True)
    next_population = []
    for i in range(topk):
        next_population.append(all_population[index_list[i]])

    return next_population



def set_seed(seed):
    np.random.seed(seed)
    random.seed(seed)


# 生成随机解的结果


def genetic_algorithm(checkpoint=10):
    population = [np.random.permutation(custom_list).copy() for _ in range(POPULATION_SIZE)]
    best_solution = max(population, key=fitness_function)
    best_fitness = fitness_function(best_solution)
    checkpoint_costs = [(0, 1 / best_fitness)]  # 初始成本，迭代次数为0

    for genetic_num in range(1, N_GENERATIONS + 1):
        fitness_values = [fitness_function(x) for x in population]
        next_population = selection(population, fitness_values)
        next_population = mutation(next_population)
        population = next_population

        current_best_solution = max(population, key=fitness_function)
        current_best_cost = 1 / fitness_function(current_best_solution)
        if fitness_function(current_best_solution) > best_fitness:
            best_solution = current_best_solution
            best_fitness = fitness_function(best_solution)

        if genetic_num % checkpoint == 0:
            checkpoint_costs.append((genetic_num, current_best_cost))
            print(f"Checkpoint at generation {genetic_num}: Best Cost = {current_best_cost}")

    return best_solution, best_fitness, checkpoint_costs

# 绘制随机解的成本分布图
def plot_routes(routes, coord, total_cost, demand, filename):
    fig, ax = plt.subplots()
    # Depot点
    ax.scatter(coord[0][0], coord[1][0], s=100, c='red')  # 红色代表仓库
    # 客户点
    ax.scatter(coord[0][1:], coord[1][1:], s=demand[1:], c='blue')  # 蓝色代表客户

    # 绘制路线
    for i, route in enumerate(routes):
        # 创建一个随机颜色
        color = np.random.rand(3, )
        plot_list = [0] + route + [0]  # 从仓库出发，到仓库结束
        ax.plot(coord[0][plot_list], coord[1][plot_list], linewidth=1, color=color)

        # 路径上的点
        ax.scatter(coord[0][plot_list], coord[1][plot_list], s=demand[plot_list], c=color,
                   label=f'路线{i + 1}:{sum(demand[route])}')

    # 设置标题和图例
    ax.set_title(f'路线数量:{len(routes)}, 总成本:{int(total_cost)}')
    ax.legend(loc=1, fontsize='x-small')
    plt.show()
    fig.savefig(filename)  # 保存图像文件




def generate_random_solutions(num_samples):
    random_costs = []
    for _ in range(num_samples):
        random_solution = np.random.permutation(custom_list)
        random_routes = solution_routes(random_solution)
        cost = routes_cost(random_routes)
        random_costs.append(cost)
    return random_costs
def plot_optimization_checkpoint(checkpoint_costs):
    # 分离数据点为迭代次数和成本
    generations, costs = zip(*checkpoint_costs)

    plt.figure(figsize=(10, 5))
    plt.plot(generations, costs, marker='o', linestyle='-', color='b')
    plt.xlabel('迭代代数')
    plt.ylabel('成本')
    plt.title('遗传算法优化过程')
    plt.xticks(generations)  # 显示所有检查点的迭代次数
    plt.grid(True)
    plt.savefig(f'../static/checkpoint_img/{int(time.time())}optimization_process_checkpoint.png')  # 保存图像文件
    plt.show()

# 运行遗传算法
set_seed(0)

# 迭代开始
# def genetic_algorithm():
#     population = [np.random.permutation(custom_list).copy() for _ in range(POPULATION_SIZE)]
#     best_solution = max(population, key=fitness_function)
#
#     for genetic_num in range(N_GENERATIONS):
#         fitness_values = [fitness_function(x) for x in population]
#
#         # print(fitness_values)
#         next_population = selection(population, fitness_values)
#         # print(next_population)
#         next_population = mutation(next_population)
#         population = next_population
#         current_best_solution = max(population, key=fitness_function)
#         if fitness_function(current_best_solution) > fitness_function(best_solution):
#             best_solution = current_best_solution
#         print(genetic_num, 1 / fitness_function(best_solution))
#
#     return best_solution, fitness_function(best_solution)


# 固定随机种子
# # 生成并绘制随机解的成本
# random_costs = generate_random_solutions(1)
# single_random_cost = random_costs[0]
#
# random_solution = np.random.permutation(custom_list)
# random_routes = solution_routes(random_solution)
#
# plot_routes(random_routes, coord, single_random_cost, demand,
#             f'../static/random_image/{int(time.time())}random_solution.png')
#
# best_solution, _ = genetic_algorithm()
# print("best_solution", best_solution)
# routes = solution_routes(best_solution)
# print("routes", routes)
# cost = routes_cost(routes)
# print('current:', cost)
# plot_routes(routes, coord, cost, demand, f'../static/optimization_image/{int(time.time())}optimization_solution.png')

#
# fig, ax = plt.subplots()
# ax.scatter(coord[0][0], coord[1][0], s=max(demand) * 2, c='r')
# ax.scatter(coord[0][1:], coord[1][1:], s=demand[1:])
# for i, route in enumerate(routes):
#     plot_list = [0] + route + [0]
#     ax.plot(coord[0][plot_list], coord[1][plot_list], linewidth='0.3',
#             label='r' + str(i + 1) + ':' + str(sum(demand[plot_list])))
# ax.legend(loc=1, fontsize='6')
# ax.set_title('route num:{}, total cost:{:.0f}'.format(len(routes), cost))
# plt.show()
# fig.savefig('../static/solution.jpg')
# routes = vrp.read_solution(r'data\A-n32-k5.sol')['routes']
# print('best:', routes_cost(routes))

def main():
    # 生成随机解并绘图
    random_solution = np.random.permutation(custom_list)
    random_routes = solution_routes(random_solution)
    random_cost = routes_cost(random_routes)
    plot_routes(random_routes, coord, random_cost, demand,
                f'../static/random_image/{int(time.time())}random_solution.png')

    # 执行遗传算法并记录优化过程
    best_solution, best_fitness, checkpoint_costs = genetic_algorithm(checkpoint=10)
    best_routes = solution_routes(best_solution)
    best_cost = routes_cost(best_routes)

    # 绘制优化结果
    plot_routes(best_routes, coord, best_cost, demand,
                f'../static/optimization_image/{int(time.time())}optimization_solution.png')

    # 绘制优化过程
    plot_optimization_checkpoint(checkpoint_costs)




main()
